# --------------- data_processing2.py ---------------
import pandas as pd
import altair as alt
import re
from typing import Optional, Tuple


class DataProcessor:
    def __init__(self, file_path: str):
        self.df = self._load_data(file_path)
        self._preprocess_data()

    def _load_data(self, file_path: str) -> pd.DataFrame:
        """加载并预处理基础数据"""
        try:
            df = pd.read_csv(file_path)
            # 添加字段重命名
            df = df.rename(columns={'workYear': 'experience'})
            df['min_salary'] = df['salary'].str.extract(r'(\d+)k-').astype(float)
            df['max_salary'] = df['salary'].str.extract(r'-(\d+)k').astype(float)
            df['avg_salary'] = (df['min_salary'] + df['max_salary']) / 2
            return df.dropna(subset=['positionName', 'city'])
        except Exception as e:
            raise ValueError(f"数据加载失败: {str(e)}")

    def _preprocess_data(self):
        """数据清洗预处理"""
        # 基础清洗
        self.df = self.df[self.df['avg_salary'].between(10, 100)]
        self.df['positionName'] = self.df['positionName'].str.replace(r'[^\w\u4e00-\u9fff]', '', regex=True)

        # 新增预处理逻辑
        self.df['education'] = self.df['education'].str.replace('不限', '无要求').fillna('无要求')
        self.df['experience'] = self.df['experience'].fillna('0-0年经验')
        self.df['experience'] = self.df['experience'].str.extract(r'(\d+-\d+年经验)').fillna('0-0年经验')

    # ----------------- 核心查询方法（共5个）-----------------
    def get_salary_analysis(self, position: str, city: str = None) -> Optional[Tuple[str, alt.Chart]]:
        """薪资分析"""
        filtered = self._filter_data(position, city)
        if len(filtered) < 5:
            return None

        base = alt.Chart(filtered).properties(width=600, height=400)
        chart = base.mark_boxplot(size=30).encode(
            x=alt.X('city:N', title='城市'),
            y=alt.Y('avg_salary:Q', title='平均薪资 (k)'),
            color=alt.Color('city:N', legend=None)
        ) + base.mark_text(dy=-10, color='red').encode(
            text='count(positionName):N'
        )

        stats = filtered.groupby('city')['avg_salary'].agg(['mean', 'count']).reset_index()
        return f"{position}岗位薪资分布（标注为样本数）", chart

    def get_city_demand(self, city: str, top: int = 5) -> Optional[Tuple[str, alt.Chart]]:
        """城市岗位需求分析"""
        city_data = self.df[self.df['city'] == city]
        if len(city_data) < 10:
            return None

        # 关键修改：显式命名列
        rank = city_data['positionName'].value_counts().head(top).reset_index()
        rank.columns = ['positionName', 'count']  # 手动指定列名

        chart = alt.Chart(rank).mark_bar().encode(
            y=alt.Y('positionName:N', title='岗位名称', sort='-x'),
            x=alt.X('count:Q', title='职位数量'),
            tooltip=['positionName:N', 'count:Q']
        ).properties(width=800, title=f"{city}热门岗位TOP{top}")

        return f"{city}市需求最大的岗位是{rank['positionName'].iloc[0]}（{rank['count'].iloc[0]}个职位）", chart

    def get_skill_analysis(self, skill: str) -> Optional[Tuple[str, alt.Chart]]:
        """技能需求分析"""
        filtered = self.df[self.df['skillLables'].apply(
            lambda x: skill.lower() in str(x).lower())]
        if len(filtered) < 10:
            return None

        chart = alt.Chart(filtered).mark_circle(size=60).encode(
            x=alt.X('companySize:N', title='公司规模'),
            y=alt.Y('avg_salary:Q', title='平均薪资'),
            color=alt.Color('education:N', title='学历要求'),
            tooltip=['positionName', 'companyFullName']
        ).interactive().properties(title=f"{skill}技能相关岗位分布")

        return f"掌握{skill}技能的岗位平均薪资为{filtered['avg_salary'].mean():.1f}k", chart

    def get_education_analysis(self, city: str = None) -> Optional[Tuple[str, alt.Chart]]:
        """学历薪资分析（新增方法）"""
        filtered = self.df if not city else self.df[self.df['city'] == city]
        if len(filtered) < 10:
            return None

        stats = filtered.groupby('education').agg(
            avg_salary=('avg_salary', 'mean'),
            count=('education', 'count')
        ).reset_index()

        chart = alt.Chart(stats).mark_bar().encode(
            x=alt.X('education:N', title='学历要求'),
            y=alt.Y('avg_salary:Q', title='平均薪资 (k)'),
            color=alt.Color('education:N', legend=None),
            tooltip=['education', 'avg_salary', 'count']
        ).properties(width=500, title=f"{city + ' ' if city else ''}学历与薪资关系")

        return "不同学历要求的薪资对比", chart

    def get_experience_analysis(self, position: str) -> Optional[Tuple[str, alt.Chart]]:
        """经验要求分析（新增方法）"""
        filtered = self.df[self.df['positionName'].str.contains(position, case=False)]
        if len(filtered) < 10:
            return None

        pattern = r'(\d+)-(\d+)年经验'
        exp_data = filtered['experience'].str.extract(pattern).astype(float)
        filtered['min_exp'] = exp_data[0]
        filtered['max_exp'] = exp_data[1]

        chart = alt.Chart(filtered).mark_rect().encode(
            x=alt.X('min_exp:Q', bin=alt.Bin(maxbins=5), title='最低经验要求'),
            y=alt.Y('max_exp:Q', bin=alt.Bin(maxbins=5), title='最高经验要求'),
            color=alt.Color('count():Q', legend=alt.Legend(title='岗位数量'))
        ).properties(title=f"{position}岗位经验要求分布")

        return f"{position}岗位常见经验要求范围", chart

    def _filter_data(self, position: str, city: str = None) -> pd.DataFrame:
        """通用数据过滤方法"""
        condition = self.df['positionName'].str.contains(position, case=False)
        if city:
            condition &= (self.df['city'] == city)
        return self.df[condition]